library(Matrix)

args=commandArgs(T)
GCMatFile="/data1/home/zhoulab/13.Mouse-Ana/SC32/SC32.GCMat.RData" #GCMat(must contain BarXCRFile's barcode!)
BarXCRFile="/data1/home/zhoulab/13.Mouse-Ana/SC32/SC32.BarTCRInfo.xls" #Bar XCR info Table!
XCRflag=1 #1,TCR;2,BCR
GeneListFile="/data1/home/zhoulab/10.code/singlecell/GeneExpTable/MouseGeneList-0927.xls"  #if "F",then use the gene list!
BarClusFile="/data1/home/zhoulab/13.Mouse-Ana/SC32/SC32.BarCluster.xls"  #if "F",then use the barcode only BarXCRFile! File Format:barcode cluster
BarMode=2 #1,use XCR Table barcode; 2, use BarCluster File barcode
CITERData=F #if not, fill "F". AllSCITEBarExpMat must in RData

OutGeneExpFile= "/data1/home/zhoulab/13.Mouse-Ana/SC32/SC32.TCRBarGeneExp.xls"

#test:
#setwd("/data4/dongxj/P68/last")
#GCMatFile="P68.GCMat.RData"
#BarXCRFile="P68.BarTCRInfo.xls"
#XCRflag="1"
#GeneListFile="HumanGeneList-0927.xls"
#BarClusFile = "P68.BarCluster.xls"
#BarMode = "2"
#CITERData = "P68.CITEExp.Mat.RData"
#OutGeneExpFile="P68.GeneExp.Test.xls"

#finish in 2021.11.11

#combine Barcode's GCMat,XCR Info Table,Cluster Table,..., create a new table!
#new table:UMI,GeneExpNum,XCRInfo,GeneExpProfile,... and so on
#if TCR's barcode not in CDNA GCMat, it will add in getGCMat.R

#read files
load(GCMatFile)

#
#BarXCRTable
readTableToDF <- function(TableFile)
{
  
  BarXCRTable <- read.table(BarXCRFile,sep = "\t",header = FALSE,stringsAsFactors = FALSE)
  HeadLine <- as.character(BarXCRTable[1,])
  colnames(BarXCRTable) <- HeadLine
  if(dim(BarXCRTable)[1] == 1)
  {
    BarXCRTable <- BarXCRTable[-1,]
  }else
  {
    BarXCRTable <- BarXCRTable[2:dim(BarXCRTable)[1],]
  }
  
  return(BarXCRTable)
}

BarXCRTable <- readTableToDF(BarXCRFile)
BarXCRTableBarcode <- BarXCRTable$barcode
#BarXCRTable end.

if(length(BarXCRTableBarcode) == 0)
{
  print("ERROR:BarXCRTable don't have barcode!")
  return(FALSE)
}

#geneList
if(GeneListFile != "F")
{
  GeneList1 <- readLines(GeneListFile)
  GCMatGene <- rownames(GCMat)
  GeneList <- intersect(GeneList1,GCMatGene)
}else
{
  GeneList <- c("Cd4","Cd8a","Cd8b","Cd79a","Cd79b","Cd19")
}
#geneList end.

#read barcode - cluster table
if(BarClusFile != "F")
{
  #BarClusFile head must be barcode - cluster!

  #cluster name add prifix "c"
  BarClusTable <- read.table(BarClusFile,sep = "\t",header = TRUE,stringsAsFactors = FALSE)
  BarClus <- paste("c",BarClusTable$cluster,sep="")
  BarClusTable$cluster <- BarClus
  rownames(BarClusTable) <- BarClusTable$barcode

}else
{
  BarClusTable <- matrix("-",ncol = 1,nrow = length(BarXCRTableBarcode))
  BarClusTable <- as.data.frame(cbind(BarXCRTableBarcode,BarClusTable))
  names(BarClusTable) <- c("barcode","cluster")
}

#read files end.

#check barcode (XCRTable barcode must all in GCMat)
GCMatBarcode <- colnames(GCMat)
commonBar <- intersect(BarXCRTableBarcode,GCMatBarcode)
if(length(commonBar) != length(BarXCRTableBarcode))
{
  NoHaveBar <- setdiff(BarXCRTableBarcode,commonBar)
  stri <- paste("WARNING:barcode error! barcode not in GCMat number:",length(NoHaveBar),sep="")
  print(stri)
  Addmat <- matrix(0,nrow=dim(GCMat)[1],ncol=length(NoHaveBar))
  rownames(Addmat) <- rownames(GCMat)
  colnames(Addmat) <- NoHaveBar
  GCMat <- cbind(GCMat,Addmat)
  commonBar <- intersect(BarXCRTableBarcode,GCMatBarcode)
}
#check barcode end!

#base function
AddGeneExp <- function(TempTable,GeneList,GCMat)
{
  TableBar <- as.character(TempTable$barcode)
  ChoMat <- as.matrix(t(GCMat[GeneList,TableBar]))
  UMI <- colSums(GCMat[,TableBar])
  GeneExpNum <- colSums(GCMat[,TableBar] > 0)

  TempTable <- cbind(TempTable,UMI)
  TempTable <- cbind(TempTable,GeneExpNum)
  TempTable <- cbind(TempTable,ChoMat)  

  return(TempTable)
}


TCRSta <- function(BarXCRTable,GCMat,commonBar)
{
#BarXCRTable col:TCR/BCR
    #test input:
    #test input end.
    
    #get temp Gene Mat
    TempGene <- c("Cd4","Cd8a","Cd79a","Cd79b","Cd19")
    ChoMat <- as.matrix(t(GCMat[TempGene,commonBar]))
    NewTable <- cbind(BarXCRTable,ChoMat)
 
    #TCR CD4 or CD8
    AllTCR <- unique(NewTable$TCRName)
    
    TCRChar1List <- c()
    for(i in AllTCR)
    {
      ThisTCRTable <- NewTable[which(NewTable$TCRName == i),]
      Cd8ExpNum = sum(ThisTCRTable$Cd8a > 0)
      Cd4ExpNum = sum(ThisTCRTable$Cd4 > 0)
  
      if(Cd4ExpNum > Cd8ExpNum)
      {
        TCRChar1 = "Cd4TCR"
      }else if(Cd8ExpNum > Cd4ExpNum)
      {
        TCRChar1 = "Cd8TCR"
      }else
      {
        TCRChar1 = "unknow"
      }
      TCRChar1List <- c(TCRChar1List,TCRChar1)
    }
    names(TCRChar1List) <- AllTCR
  
    TCRChar1List2 <- c()
    for(i in commonBar)
    {
      ThisBarTCR <- NewTable[which(NewTable$barcode == i),]$TCRName[1]
      TCRChar1List2 <- c(TCRChar1List2,as.character(TCRChar1List[ThisBarTCR]))
    }
    #TCR CD4 or CD8 end.
    
    #TCR tumor special
    TCRChar3List <- c()  #every TCR tumor special info
    
    TableHead <- colnames(NewTable)
    if(c("Tumor") %in% TableHead & c("Normal") %in% TableHead & c("Blood") %in% TableHead)
    {
        tflag <- 1
    }else if(c("Tumor") %in% TableHead & c("Normal") %in% TableHead & !c("Blood") %in% TableHead)
    {
        tflag <- 2
    }else if(c("Tumor") %in% TableHead & !c("Normal") %in% TableHead & c("Blood") %in% TableHead)
    {
        tflag <- 3
    }else
    {
        tflag <- 0
    }

    
    if(tflag > 0)
    {
        for(i in AllTCR)
        {

            ThisTCRTable <- NewTable[which(NewTable$TCRName == i),]
      
            TumorNum <- as.integer(ThisTCRTable$Tumor[1])
            if(tflag == 1)
            {
                ControlNum <- as.integer(ThisTCRTable$Normal[1]) + as.integer(ThisTCRTable$Blood[1])
            }else if(tflag == 2)
            {
                ControlNum <- as.integer(ThisTCRTable$Normal[1])
            }else if(tflag == 3)
            {
                ControlNum <- as.integer(ThisTCRTable$Blood[1])
            }else
            {
                ControlNum <- TumorNum
                print("error in TCR tumor special!")
            }

            if(TumorNum > 5 & ControlNum < 3)
            { 
                TCRChar3 <- "TumorSpecial"
            }else
            {
                TCRChar3 <- "-"
            }

            TCRChar3List <- c(TCRChar3List,TCRChar3)
        }   
 
        #final view 
        names(TCRChar3List) <- AllTCR
        TCRChar3List2 <- c() #final View TCRinfo in Table
        for(i in commonBar)
        {
          ThisBarTCR <- NewTable[which(NewTable$barcode == i),]$TCRName[1]
          TCRChar3List2 <- c(TCRChar3List2,as.character(TCRChar3List[ThisBarTCR]))
        }
    }else
    {
        TCRChar3List2 <- rep("-",time = length(commonBar))
    }
    #TCR tumor special end.

    #TCR doublet
    TCRChar2List <- c()
    for(i in commonBar)
    {
      ThisTCRTable <- NewTable[which(NewTable$barcode == i),]
      Cd4Expflag <- 0
      Cd8Expflag <- 0
      BMExpflag <- 0
      
      if(ThisTCRTable$Cd4 > 0)
      {
        Cd4Expflag = 1
      }
      
      if(ThisTCRTable$Cd8a > 0)
      {
        Cd8Expflag = 1
      }
      
      if(ThisTCRTable$Cd79a > 0 | ThisTCRTable$Cd79b > 0 | ThisTCRTable$Cd19 > 0)
      {
        BMExpflag = 1
      }
      
      if(Cd4Expflag == 1 & Cd8Expflag == 1)
      {
        TCRChar2 <- "doublet"
      }else if(BMExpflag == 1)
      {
        TCRChar2 <- "doublet"
      }else
      {
        TCRChar2 <- "-"
      }
      TCRChar2List <- c(TCRChar2List,TCRChar2)
    }
    #TCR doublet end.

    #add these col in Table
    RawHead <- colnames(BarXCRTable)
    NewHead <- c(RawHead,"TCRatt1","TCRatt2","TCRatt3")
    NewTable2 <- cbind(BarXCRTable,TCRChar1List2,TCRChar2List,TCRChar3List2)
    colnames(NewTable2) <- NewHead
    #add end.

    return(NewTable2)
}

DataFramAddNull <- function(PreBar,TempTable)
{
  BarTempTable <- TempTable$barcode
  ThisCommBar <- intersect(PreBar,BarTempTable)
  ATable <- BarTempTable[ThisCommBar,]
  RemainBar <- setdiff(PreBar,ThisCommBar)

  if(length(RemainBar) > 0)
  {
    BTable <- as.data.frame(cbind(RemainBar,matrix("-",nrow = RemainBar,ncol=dim(TempTable)[2]-1)))
    FTable <- rbind(ATable,BTable)
  }else
  {
    FTable <- ATable
  }

  return(FTable)
}


FindTableInfo <- function(keyList,Table)
{

  finalTable <- NULL
  FillNULL <- rep("-",time=dim(Table)[2])
  
  TableKey <- Table[,1]
  for(i in keyList)
  {
    if(i %in% TableKey)
    {
      finalTable <- rbind(finalTable,as.matrix(Table[which(Table$barcode == i),][1,]))
    }else
    {
      finalTable <- rbind(finalTable,FillNULL)
    }
    
  }
  
  
  finalTable <- finalTable[,2:dim(finalTable)[2]]
  finalTable <- as.data.frame(finalTable)
  
  return(finalTable)
}

TabAddClus <- function(XCRTable,ClusTable)
{
  #ClusTable must have tow col:barcode,cluster,and has rownames
  AddClusTab <- rep("-",length(XCRTable$barcode))
  names(AddClusTab) <- XCRTable$barcode

  CommBars <- intersect(names(AddClusTab),ClusTable$barcode)
  AddClusTab[CommBars] <- ClusTable[CommBars,]$cluster
 
  XCRTableHead <- colnames(XCRTable)
  Head <- c("barcode","cluster",XCRTableHead[2:length(XCRTableHead)])
  FinalTab <- as.data.frame(XCRTable$barcode)
  FinalTab <- cbind(FinalTab,AddClusTab)
  FinalTab <- cbind(FinalTab,XCRTable[,2:dim(XCRTable)[2]])
  colnames(FinalTab) <- Head
  rownames(FinalTab) <- XCRTable$barcode

  return(FinalTab)
}
#base function end.

#main

#add TCR att1/att2
if(XCRflag == "1")  #TCR
{
  print("raw table")
  NewTable2 <- TCRSta(BarXCRTable,GCMat,commonBar)
}else #BCR
{
  NewTable2 <- BarXCRTable
}

#add cluster info
print("add cluster info")
NewTable3 <- TabAddClus(NewTable2,BarClusTable)
  
#add GeneExp
print("add view genes")
NewTable4 <- AddGeneExp(NewTable3,GeneList,GCMat)
write.table(NewTable4,OutGeneExpFile,sep="\t",quote=F,row.names = FALSE,col.names = TRUE)

# #add CITESEQ table
# if(CITERData != "F")
# {
#   print("Table need add CITESEQ part")
#   TabBars <- rownames(NewTable4)
#   citeflag <- 0
#   load(CITERData)
#   CITEExpMat <- CITEExpMat[TabBars,]  
# 
#   allcites <- colnames(CITEExpMat)
#   CiteName <- allcites[grep("^CITE-",allcites)]
#   ProtName <- allcites[grep("^Protein_",allcites)]
#   
#   if(length(CiteName) > 0 & length(ProtName) == 0)
#   {
#     print("only CITESeq no Protein.")
#     citeflag <- 1
#     CITEExpMat1 <- CITEExpMat
#     
#   }else if(length(CiteName) > 0 & length(ProtName) > 0)
#   {
#     print("CITESeq + protein.")
#     #protein match mrna
#     citeflag <- 2
# 
#     AllMrnas <- rownames(GCMat)
# 
#     NameIndex <- colnames(CITEExpMat)[1:length(CiteName)]  #like:CITE-1,CITE-2,Protein_GA,MRna_GA,...    
#     NeedMrnaGenes <- c()
#     for(i in ProtName)
#     {
#       ThisPro <- sub("^Protein_","",i)
#       if(ThisPro %in% AllMrnas)
#       {
#         NeedMrnaGenes <- c(NeedMrnaGenes,ThisPro)
# 
#         ThisMrnaName <- paste("MRna_",ThisPro,sep="")
#         NameIndex <- c(NameIndex,i)
#         NameIndex <- c(NameIndex,ThisMrnaName)
#       }else
#       {
#         NameIndex <- c(NameIndex,i)
#       }
#     }
# 
#     CITEBars <- rownames(CITEExpMat)    
#     NeedMrnaMat <- GCMat[NeedMrnaGenes,CITEBars]
#     rownames(NeedMrnaMat) <- paste("MRna_",rownames(NeedMrnaMat),sep="")
#     NeedMrnaMat <- t(NeedMrnaMat)
#     
#     CITEExpMat <- as.matrix(CITEExpMat)
#     NeedMrnaMat <- as.matrix(NeedMrnaMat)
#     print(NeedMrnaMat[1:5,1:5])
#     print(CITEExpMat[1:5,1:5])
#     CITEExpMatM <- cbind(CITEExpMat,NeedMrnaMat)
#     CITEExpMatM <- CITEExpMatM[,NameIndex]
#   }
#   
#   #if table barcodes more than CITESEQ matrix, will fill 0
#   #TableBar <- as.character(NewTable4$barcode)  
#   #TableAndCITECommBar <- intersect(TableBar,rownames(CITEExpMat1))
#   #if(length(TableAndCITECommBar) != length(TableBar))
#   #{
#   #  print("warning:Gene Exp Table barcodes not equal CITESeq barcodes! we will choose common part as new table!")
#   #  
#   #  NeedAddCITEBar <- setdiff(TableBar,TableAndCITECommBar)
#   #  if(length(NeedAddCITEBar) > 0)
#   #  {
#   #    tempMat <- matrix(0,nrow = length(NeedAddCITEBar),ncol = dim(CITEExpMat1)[2])
#   #    rownames(tempMat) <- NeedAddCITEBar
#   #    colnames(tempMat) <- colnames(CITEExpMat1)
#   #    tempMat<- as(tempMat,"dgCMatrix")
#   #    
#   #    CITEExpMat1 <- rbind(CITEExpMat1,tempMat)
#   #  }
#   #}
#   #if table barcodes more than CITESEQ matrix, will fill 0
#   
#   #
#   Head <- colnames(NewTable4)
#   Head <- c(Head,colnames(CITEExpMatM))
#   NewTable5 <- cbind(NewTable4,CITEExpMatM)
#   colnames(NewTable5) <- Head
#   
# }else
# {
#   #no CITESEQ
#   NewTable5 <- NewTable4
# }
# #add CITESEQ table
# 
# #output
# write.table(NewTable5,OutGeneExpFile,sep="\t",quote=F,row.names = FALSE,col.names = TRUE)
